package com.example.smartwaterapi.filter;

import com.alibaba.fastjson2.JSON;
import com.example.smartwaterapi.common.api.RestApiResponse;
import jakarta.servlet.FilterChain;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Component;
import org.springframework.web.filter.OncePerRequestFilter;

import java.io.IOException;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * IP限流过滤器
 */
@Slf4j
@Component
public class IpRateLimitFilter extends OncePerRequestFilter {
    
    @Autowired
    private RedisTemplate<String, Object> redisTemplate;
    
    @Value("${rate-limit.ip.enable:true}")
    private boolean ipLimitEnable;
    
    @Value("${rate-limit.ip.max-requests:800}")
    private int ipMaxRequests;
    
    @Value("${rate-limit.ip.window-seconds:60}")
    private int ipWindowSeconds;
    
    @Value("${rate-limit.ip.strict-max-requests:30}")
    private int strictMaxRequests;
    
    @Value("#{'${rate-limit.ip.whitelist:}'.split(',')}")
    private List<String> ipWhitelist;
    
    @Value("#{'${rate-limit.ip.strict-paths:}'.split(',')}")
    private List<String> strictPaths;
    
    @Override
    protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, 
                                    FilterChain filterChain) throws ServletException, IOException {
        
        // 如果未启用IP限流，直接放行
        if (!ipLimitEnable) {
            filterChain.doFilter(request, response);
            return;
        }
        
        // 获取客户端IP地址
        String clientIp = getClientIp(request);
        String uri = request.getRequestURI();
        
        // 对静态资源、Swagger等不进行限流
        if (isResourceRequest(uri)) {
            filterChain.doFilter(request, response);
            return;
        }
        
        // 检查IP白名单
        if (isWhitelistedIp(clientIp)) {
            log.debug("IP {} 在白名单中，跳过限流", clientIp);
            filterChain.doFilter(request, response);
            return;
        }
        
        // 确定限流策略
        boolean isStrictPath = isStrictPath(uri);
        int maxRequests = isStrictPath ? strictMaxRequests : ipMaxRequests;
        String limitType = isStrictPath ? "strict" : "normal";
        
        // 限流键名
        String redisKey = String.format("rate_limit:ip:%s:%s", limitType, clientIp);
        
        // 使用滑动窗口算法进行更精确的限流
        if (isRateLimitExceeded(redisKey, maxRequests)) {
            handleIpRateLimit(response, clientIp, uri, maxRequests, limitType);
            return;
        }
        
        // 记录访问
        recordIpAccess(redisKey);
        
        // 继续处理请求
        filterChain.doFilter(request, response);
    }
    
    /**
     * 获取客户端IP地址
     */
    private String getClientIp(HttpServletRequest request) {
        String ip = request.getHeader("X-Forwarded-For");
        
        if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("Proxy-Client-IP");
        }
        if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("WL-Proxy-Client-IP");
        }
        if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("HTTP_CLIENT_IP");
        }
        if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("HTTP_X_FORWARDED_FOR");
        }
        if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getRemoteAddr();
        }
        
        // 如果是多级代理，取第一个IP
        if (ip != null && ip.contains(",")) {
            ip = ip.split(",")[0].trim();
        }
        
        return ip;
    }
    
    /**
     * 判断是否是资源请求（静态资源、Swagger等）
     */
    private boolean isResourceRequest(String uri) {
        return uri.contains("/swagger") 
                || uri.contains("/v3/api-docs") 
                || uri.contains("/doc.html")
                || uri.contains("/favicon.ico")
                || uri.contains("/static/")
                || uri.contains("/actuator/health")
                || uri.endsWith(".js")
                || uri.endsWith(".css")
                || uri.endsWith(".png")
                || uri.endsWith(".jpg")
                || uri.endsWith(".jpeg")
                || uri.endsWith(".gif")
                || uri.endsWith(".svg");
    }
    
    /**
     * 检查IP是否在白名单中
     */
    private boolean isWhitelistedIp(String clientIp) {
        if (ipWhitelist == null || ipWhitelist.isEmpty()) {
            return false;
        }
        
        for (String whitelistPattern : ipWhitelist) {
            if (whitelistPattern.trim().isEmpty()) {
                continue;
            }
            
            // 支持通配符匹配
            if (whitelistPattern.contains("*")) {
                String regex = whitelistPattern.replace(".", "\\.")
                                              .replace("*", ".*");
                if (clientIp.matches(regex)) {
                    return true;
                }
            } else if (clientIp.equals(whitelistPattern.trim())) {
                return true;
            }
        }
        
        return false;
    }
    
    /**
     * 检查是否是严格限流路径（主要针对管理员接口等特殊接口）
     */
    private boolean isStrictPath(String uri) {
        if (strictPaths == null || strictPaths.isEmpty()) {
            return false;
        }
        
        for (String strictPath : strictPaths) {
            if (strictPath.trim().isEmpty()) {
                continue;
            }
            if (uri.contains(strictPath.trim())) {
                return true;
            }
        }
        
        return false;
    }
    
    /**
     * 使用滑动窗口算法检查是否超过限流
     */
    private boolean isRateLimitExceeded(String redisKey, int maxRequests) {
        try {
            long now = System.currentTimeMillis();
            long windowStart = now - ipWindowSeconds * 1000L;
            
            // 移除过期的访问记录
            redisTemplate.opsForZSet().removeRangeByScore(redisKey, 0, windowStart);
            
            // 获取当前窗口内的访问次数
            Long count = redisTemplate.opsForZSet().zCard(redisKey);
            int currentCount = (count != null) ? count.intValue() : 0;
            
            log.debug("IP限流检查: key={}, 当前访问次数: {}/{}", redisKey, currentCount, maxRequests);
            
            return currentCount >= maxRequests;
            
        } catch (Exception e) {
            log.error("检查IP限流状态失败: {}", e.getMessage(), e);
            // 发生异常时不限流，避免影响正常业务
            return false;
        }
    }
    
    /**
     * 记录IP访问
     */
    private void recordIpAccess(String redisKey) {
        try {
            long now = System.currentTimeMillis();
            
            // 使用时间戳作为分数和值，确保唯一性
            redisTemplate.opsForZSet().add(redisKey, 
                String.valueOf(now + "_" + Thread.currentThread().getId()), now);
            
            // 设置过期时间，避免内存泄漏
            redisTemplate.expire(redisKey, ipWindowSeconds * 2, TimeUnit.SECONDS);
            
        } catch (Exception e) {
            log.error("记录IP访问失败: {}", e.getMessage(), e);
        }
    }
    
    /**
     * 处理IP限流响应
     */
    private void handleIpRateLimit(HttpServletResponse response, String clientIp, 
                                  String uri, int maxRequests, String limitType) throws IOException {
        
        log.warn("IP {} 访问 {} 被限流，类型: {}, 限制: {}/{}秒", 
                clientIp, uri, limitType, maxRequests, ipWindowSeconds);
        
        response.setContentType(MediaType.APPLICATION_JSON_VALUE);
        response.setCharacterEncoding("UTF-8");
        response.setStatus(429); // 429 Too Many Requests
        
        // 设置响应头
        response.setHeader("Retry-After", String.valueOf(ipWindowSeconds));
        response.setHeader("X-RateLimit-Limit", String.valueOf(maxRequests));
        response.setHeader("X-RateLimit-Window", String.valueOf(ipWindowSeconds));
        response.setHeader("X-RateLimit-Type", limitType);
        
        String message;
        if ("strict".equals(limitType)) {
            message = String.format("敏感接口访问过于频繁，请 %d 秒后重试", ipWindowSeconds);
        } else {
            message = String.format("请求过于频繁，请 %d 秒后重试", ipWindowSeconds);
        }
        
        RestApiResponse<String> errorResponse = RestApiResponse.custom(
                org.springframework.http.HttpStatus.TOO_MANY_REQUESTS, 
                "429", 
                message,
                null);
        
        response.getWriter().write(JSON.toJSONString(errorResponse));
        response.getWriter().flush();
    }
} 